import pandas as pd
import openpyxl
import sys


def strip_space(df):
    df = df.rename(columns={column_name: column_name.strip() for column_name in df.columns})
    df = df.applymap(lambda x: x.strip() if isinstance(x, str) else x)
    return df

def read_data_wx(path) -> pd.DataFrame: 
    df = pd.read_csv(path, header=16, encoding='utf-8')
    # print(data)
    df = df.iloc[:, [0, 4, 7, 2, 3, 5]]  # ingest columns in order
    df = strip_space(df)
    df.iloc[:, 0] = df.iloc[:, 0].astype('datetime64')
    df.iloc[:, 5] = df.iloc[:, 5].replace(regex={r'^.': ''}).astype('float64')  # amount float type
    df.rename(columns={'当前状态': '支付状态', '交易类型': '类型', '金额(元)': '金额'}, inplace=True)
    df.insert(1, '来源', "微信")
    print("成功读取 " + str(len(df)) + " 条「微信」账单数据")
    return df

def read_data_alipay(path) -> pd.DataFrame: 
    df = pd.read_csv(path, header=4, skipfooter=7, encoding='gbk', engine='python')
    # print(data)
    df = df.iloc[:, [2, 10, 11, 7, 8, 9]]  # ingest columns in order
    df = strip_space(df)
    df.iloc[:, 0] = df.iloc[:, 0].astype('datetime64')
    df.iloc[:, 5] = df.iloc[:, 5].astype('float64')  # amount float type
    df = df[~df['收/支'].str.contains('其他')]  
    df.rename(columns={'交易创建时间': '交易时间', '交易状态': '支付状态', '商品名称': '商品', '金额（元）': '金额'}, inplace=True)
    df.insert(1, '来源', "支付宝")
    print("成功读取 " + str(len(df)) + " 条「支付宝」账单数据")
    return df

def classify(df):
    df.insert(6, '分类', '')
    # 简餐 meal
    meal = df['商品'].str.contains( '|'.join([
        '餐', '食', '吃', '饭', '锅' ,'包', '汤', '面', '饺', '寿司', '便利', '外卖', '美团', '饿了么']) )
    df.loc[meal, '分类'] = '简餐'  # 后期考虑用配置文件配置

    # 饕餮 & 社交 banquet and social
    banquet = meal & (df.loc[meal]['金额'] > 40)  # distribute the high-cost meal
    df.loc[banquet, '分类'] = '饕餮&社交'

    # 零食 snack
    df_snack = df.stack().str.contains('|'.join([
        'coff', 'cafe', '咖啡', '茶', '小e']), case=False).unstack()
    snack = df_snack.any(axis=1)  # transform to Series
    df.loc[snack, '分类'] = '零食'  

    # 交通 transport
    df_transport = df.stack().str.contains('|'.join([
        '滴滴', '高德', '地铁', '单车', '公交', '出行', '师傅', '物业', '停车', '粤B']), case=False).unstack()
    transport = df_transport.any(axis=1)
    df.loc[transport, '分类'] = '交通'  

    # 通信&月付产品 subscription
    df_subscription = df.stack().str.contains('|'.join([
        'vip', '百度网盘', '喜马拉雅', '爱奇艺', '网易云', '充值']), case=False).unstack()
    subscription = df_subscription.any(axis=1)
    df.loc[subscription, '分类'] = '通信&月付产品'

    print('分类 [简餐] 数据 ' + str( len(df[df['分类'].isin(['简餐'])]))  + ' 条')
    print('分类 [饕餮 & 社交] 数据 ' + str(len(df.loc[banquet]))  + ' 条')
    print('分类 [零食] 数据 ' + str(len(df.loc[snack]))  + ' 条')
    print('分类 [交通] 数据 ' + str(len(df.loc[transport]))  + ' 条')
    print('分类 [通信 & 月付产品] 数据 ' + str(len(df.loc[subscription]))  + ' 条')
    print('## 未分类数据 ' + str( len(df[df['分类'].isin([''])]))  + ' 条')

    return df

    
if __name__ == '__main__':

    argv = sys.argv[1:]
    # path_wx = './test_wx.csv'
    # path_alipay = './test_alipay.csv'
    path_wx = argv[0]
    path_alipay = argv[1]
    d_wx = read_data_wx(path_wx)
    d_alipay = read_data_alipay(path_alipay)
    # merge data
    data_merge = pd.concat([d_wx, d_alipay], ignore_index=True)  
    print("总计合并数据 " + str(len(data_merge)) + " 条\n")
    data_merge = classify(data_merge)

    # merge_list = data_merge.values.tolist()  # transform DataFrame->List
    path_xlsx = 'merge.xlsx'
    data_merge.to_excel(path_xlsx, index=False)  # save
    # format xl
    workbook = openpyxl.load_workbook(path_xlsx)
    sheet = workbook['Sheet1']
    sheet.column_dimensions['F'].width = 40
    sheet.column_dimensions['E'].width = 20
    sheet.column_dimensions['A'].width = 20
    maxrow = sheet.max_row
    type_list = []
    for i in ['大件好物', '购物欲', '衣物', '学习&文具', '日用起居', '放松', '医疗', '房租', '其他', '简餐', '饕餮&社交', '零食', '交通', '通信&月付产品', '养车', '医疗']:
        type_list.append(i.split())
    for row in type_list:
        sheet.append(row)

    workbook.save(path_xlsx)
    print("write successfully!\n")
